LLAISYS 训练营结项报告

1. 项目概述

本次训练营我主要完成了 LLAISYS 的 NVIDIA 后端接入和多卡分布式推理支持，打通了构建系统、运行时、算子、模型推理链路以及多卡通信。

我最终完成的内容包括：

• 完成 Project #2：为 LLAISYS 接入 NVIDIA CUDA 后端
• 完成 Project #5：为 LLAISYS 接入基于 NCCL 的 8 卡 Tensor Parallel 推理能力，跑通了 Qwen2 在 8 张 A100 上的分布式推理流程
• 完成从 correctness 到性能分析的完整验证

2. 实验环境与复现方式

2.1 基础构建流程

每次修改 C++ / CUDA 后端后，我使用下面的流程重新构建：

xmake f --nv-gpu=y -cv
xmake
xmake install
pip install ./python

说明：Python 测试默认导入已安装的 llaisys 包，因此后端修改后必须重新执行 pip install ./python。

2.2 我实际使用的验证命令

分布式验证前，我的 shell 环境中已包含以下配置：

source /data/shared/miniconda3/etc/profile.d/conda.sh
conda activate szt
export PATH=$PATH:/usr/local/cuda/bin/:/usr/local/mpi/bin/

export NCCL_IB_HCA=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1
export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_DISABLE=0
export NCCL_IB_RETRY_CNT=7
export NCCL_IB_TIMEOUT=23
export NCCL_LAUNCH_MODE=GROUP

其中 LLAISYS_DIST_BOOTSTRAP_PATH 不需要手工设置，测试脚本内部会自动生成临时路径用于 NCCL bootstrap。

单卡 NVIDIA Runtime 验证：

python test/test_runtime.py --device nvidia

单卡 NVIDIA 推理验证：

python test/test_infer.py \
  --model /home/st3to/DeepSeek-R1-Distill-Qwen-1.5B \
  --test \
  --device nvidia

8 卡通信验证：

python test/test_dist.py --device nvidia --world-size 8

8 卡分布式推理验证：

python test/test_infer_dist.py \
  --model /home/st3to/DeepSeek-R1-Distill-Qwen-1.5B \
  --device nvidia \
  --test \
  --max_steps 128 \
  --world-size 8

3. Project #2：Integrate CUDA into LLAISYS

3.1 项目目标

Project #2 的目标是把 LLAISYS 从 CPU-only 扩展为支持 NVIDIA CUDA 推理的异构系统，主要包括：

1. 让工程能正确编译和链接 CUDA 代码
2. 让系统具备 NVIDIA Runtime API
3. 让核心算子真正走 NVIDIA 后端
4. 让 Qwen2 的整条推理链跑在 GPU 上

3.2 实现路径

第一层：构建系统

这一层我主要完成了：

• 新增 xmake/nvidia.lua
• 在 xmake.lua 中接入 --nv-gpu=y
• 新增 ENABLE_NVIDIA_API 宏控制
• 建立 llaisys-device-nvidia 和 llaisys-ops-nvidia 两个 CUDA 目标
• 链接 cudart、cublas
• 启用 CUDA device linking

初期遇到过 __cudaRegisterLinkedBinary_* 相关错误，最终定位为缺少 CUDA device linking。

第二层：NVIDIA Runtime

我对照 CPU Runtime 的接口，把 NVIDIA 对应实现补齐，主要包括：

• cudaGetDeviceCount
• cudaSetDevice
• cudaDeviceSynchronize
• cudaStreamCreate / Destroy / Synchronize
• cudaMalloc / cudaFree
• cudaMallocHost / cudaFreeHost
• cudaMemcpy / cudaMemcpyAsync

同时我抽出了统一的 CUDA 辅助层，主要负责：

• CUDA_CHECK
• stream 类型转换
• memcpy kind 转换
• fp16 / bf16 / float 的统一处理

第三层：CUDA 算子

本次我完成了以下 NVIDIA 算子：

• add
• argmax
• embedding
• linear
• rms_norm
• rope
• self_attention
• swiglu
• rearrange

实现策略如下：

• 逐元素算子如 add、swiglu：grid-stride loop
• 规约类算子如 argmax、rms_norm：shared memory reduction
• linear：cublasGemmEx
• rope、self_attention：原生 CUDA kernel
• rearrange：contiguous 快路径 + 通用 stride-aware copy

第四层：模型推理集成

在模型侧，我完成了 Qwen2 的 NVIDIA 推理支持，主要包括：

• 权重分配在 GPU
• 中间 buffer 分配在 GPU
• KV Cache 分配在 GPU
• 算子 dispatch 到 NVIDIA 实现
• 修复 linear 无 bias 路径
• 修复 Python 包装层 bias is None
• 用 rearrange 支持 KV Cache 写回

3.3 设计方案与思考

1. 为什么要先打通构建系统和 Runtime

因为 .cu 编译、device linking、stream、allocator、memcpy 这类基础能力决定了上层算子和模型是否有稳定运行环境。

2. 为什么 linear 直接选 cublasGemmEx

linear 是最核心、最频繁的计算之一，直接使用成熟的高性能库更稳妥，也更符合工程取舍。

3. 为什么低精度输入通常先转 float

像 add、rms_norm、swiglu、attention 分数计算这类路径，如果直接在低精度上累计，数值误差会更明显，因此采用“低精度存储、float 计算、低精度写回”。

4. 为什么性能优化不应该只盯着 kernel

单卡推理跑通后，主要瓶颈并不在 kernel 本身，而在固定开销：

• 每次 linear 都创建和销毁 cuBLAS handle
• Qwen2Model 每生成一个 token 都重新申请中间 buffer

所以后面我做的关键优化其实是：

• 增加可复用的 per-device cublasHandle_t
• 复用 Qwen2Model 的中间 buffer
• 保留 KV Cache contiguous 路径的纯 D2D copy

3.4 遇到的典型问题

1. CUDA 设备链接问题

初期出现过“编译通过但动态库加载失败”的情况，最后确认是缺少 CUDA device linking。这个问题如果没有完整经历过，很容易误以为是普通链接错误。

2. Python 测试读到旧包

测试导入的是已安装的 llaisys 包，因此每次都必须重新执行：

pip install ./python

3. self_attention 测试脚本本身有问题

PyTorch 参考路径里 attn_bias 在 CUDA、mask 却在 CPU，导致参考实现先报错。

3.5 测试结果

Runtime 测试

python test/test_runtime.py --device nvidia

结果：通过。

NVIDIA 算子测试

通过的算子测试包括：

• test/ops/add.py --device nvidia
• test/ops/argmax.py --device nvidia
• test/ops/embedding.py --device nvidia
• test/ops/linear.py --device nvidia
• test/ops/rms_norm.py --device nvidia
• test/ops/rope.py --device nvidia
• test/ops/self_attention.py --device nvidia
• test/ops/swiglu.py --device nvidia

端到端推理测试

python test/test_infer.py \
  --model /home/st3to/DeepSeek-R1-Distill-Qwen-1.5B \
  --test \
  --device nvidia

结果：通过，且 LLAISYS 生成的 token 与 Hugging Face / PyTorch 对齐。

3.6 性能结果与结论

下面的数据来自同一台机器、同一条单卡端到端测试路径：

python test/test_infer.py \
  --model /home/st3to/DeepSeek-R1-Distill-Qwen-1.5B \
  --test \
  --device nvidia

在这条测试路径下：

• 优化前：Hugging Face / PyTorch 约 2.57s，LLAISYS NVIDIA 约 42.94s
• 优化后：Hugging Face / PyTorch 约 2.89s，LLAISYS NVIDIA 约 1.20s

结论：

• 当前实现已完成 correctness 验证，并达到较好的性能表现
• 固定开销优化比盲目微调单个 kernel 更重要

4. Project #5：Distributed Inference

4.1 项目目标

Project #5 的目标，是在单卡 NVIDIA 推理基础上，让 LLAISYS 支持 8 卡分布式推理，并为后续更大模型和 MoE 扩展准备底座。核心问题包括：

• 多卡之间怎么通信
• 哪些算子应该切分，怎么切分
• 模型权重怎么按 rank 加载
• KV Cache 怎么按分片保存
• 最终怎么证明 8 卡输出和单卡一致

4.2 实现路径

第一层：通信层

这一层我新增了基于 NCCL 的封装：

• src/device/nvidia/nccl_context.hpp
• src/device/nvidia/nccl_context.cu

并把分布式能力接到 Context 上，支持：

• initDistributed(rank, world_size)
• allReduce
• allGather
• broadcast
• barrier

通信接口保持为通用 collective 原语，便于后续继续扩展 Tensor Parallel 和 MoE。

第二层：算子并行层

我增加了几组 Tensor Parallel 基础原语：

• columnParallelLinear
• rowParallelLinear
• gatherLastDim
• parallelEmbedding

它们对应的直觉分别是：

• 输出维切分的矩阵乘
• 输入维切分的矩阵乘并在结果处求和
• 把局部分片重新拼回完整张量
• 在词表分片上做 embedding 查询并汇总结果

第三层：模型并行层

在模型层，我对 Qwen2 做了分布式改造，主要包括：

• Python 侧按 rank 对权重切片后加载
• embed_tokens 和 lm_head 走 vocab parallel
• q/k/v、gate_proj、up_proj 走 column parallel
• o_proj、down_proj 走 row parallel
• KV Cache 只保留本 rank shard
• 最终 token 选择改成 distributed argmax

第四层：验证层

这一阶段我新增并通过了：

• test/test_dist.py
• test/test_infer_dist.py

分别用于验证 collective correctness 和 8 卡推理输出一致性。

4.3 设计方案与思考

1. 为什么采用“一进程一卡”

我采用的是“一进程绑定一张 GPU”的方式，主要原因有三点：

• 设备上下文更清晰
• communicator 生命周期更容易管理
• 和后续 torch.multiprocessing / torchrun 的组织方式更一致

2. 为什么把通信层挂在 Context

Context 原本就负责当前设备、runtime、stream、allocator 等底层能力，把 rank、world size、communicator 和 collective 入口放进去可以保持分层一致。

3. 为什么这次没有直接做理想的 head parallel

当前 Qwen-1.5B 配置里：

• num_heads = 12
• num_kv_heads = 2
• world_size = 8

这意味着 query heads 和 kv heads 都不能被 8 整除，因此无法直接做理想的 head parallel。

所以我最终选择的是更通用的方案：

• q/k/v 先按 feature 维切分
• 在进入 RoPE / attention 前 gather 回完整张量

这种方式的代价是通信更多，但有两个优势：

• 当前 Qwen-1.5B 可以在 8 卡上稳定跑通
• 先把 Tensor Parallel 的通信骨架和模型路径走通

4. 为什么要做 distributed argmax

如果在 lm_head 之后直接 gather 全量 logits，每张卡都要重新物化完整词表结果，通信量和显存开销都偏大。

所以我做了 distributed argmax：

1. 每个 rank 先做本地 argmax
2. 只收集本地最优分数和对应全局 token id
3. 最后在 host 上决定全局最优 token

5. 为什么要特别处理 non-contiguous tensor

张量并行过程中经常会出现切片视图，而这些视图不一定是 contiguous。由于 linear 的 GEMM 路径要求输入连续，因此需要显式使用 rearrange 做连续化。

4.4 遇到的关键问题

1. NCCL 与设备上下文问题

多卡通信初期遇到过 unhandled cuda error，最终排查发现是 runtime 在切换到目标 GPU 之前就创建了 stream，导致 rank 和设备上下文不匹配。

2. bootstrap 文件旧状态污染

NCCL 初始化需要共享 ncclUniqueId。共享 bootstrap 文件方式容易在第二次运行时读到旧文件，导致初始化异常。

因此我后来补了两层保护：

• rank 0 写新文件前先删除旧文件
• 非 root rank 只接受短时间窗口内的新文件

3. 分布式计时口径容易误导

一开始单卡和多卡的 load / generate 计时口径不一致，后来统一拆成 load latency 和 generate latency。

4.5 测试结果

通信层验证

python test/test_dist.py --device nvidia --world-size 8

这一部分不是只看程序能否启动，而是逐项验证基础 collective 的数学正确性。验证内容包括：

• allReduce
• allGather
• broadcast
• barrier

具体验证方式是：

• 为每个 rank 构造可区分的输入
• 调用对应 NCCL collective
• 将结果与手工构造的期望值逐一比较

这些期望值不是来自外部模型或数据集，而是直接在测试脚本中手工构造：

• allReduce：每个 rank 输入 [rank+1, ..., rank+1]，期望结果为 1 + 2 + ... + world_size
• allGather：期望结果为按 rank 顺序堆叠后的张量
• broadcast：期望结果为 root rank 的固定向量 [10, 11, 12, 13]

只有当各 rank 返回结果都与期望一致时，才认为通信层正确。

单卡推理验证

python test/test_infer.py \
  --model /home/st3to/DeepSeek-R1-Distill-Qwen-1.5B \
  --device nvidia \
  --test \
  --max_steps 4

这一步作为分布式验证的 baseline，作用是确认最新后端下单卡路径仍然正确，并提供后续分布式 token 对比的参考结果。

这里的 baseline 数据来源于单卡 LLAISYS Qwen2 在相同 prompt、相同 max_steps 下生成的 token 序列，而不是直接使用 Hugging Face 输出作为分布式验证基准。

8 卡分布式推理验证

python test/test_infer_dist.py \
  --model /home/st3to/DeepSeek-R1-Distill-Qwen-1.5B \
  --device nvidia \
  --test \
  --max_steps 4 \
  --world-size 8

这一部分的 correctness 验证重点不是只看“8 卡能跑起来”，而是验证分布式前向是否与单卡保持一致。具体验证方式是：

• 启动 8 个 worker process，并完成 NCCL 初始化
• 每个 rank 只加载自己的权重 shard
• 在相同 prompt、相同 max_steps 下执行推理
• 将 8 卡生成的最终 token 序列与单卡 baseline 逐 token 比较

因此，这里的验证数据来源是：

• 通信层：测试脚本手工构造的数学期望值
• 推理层：单卡 LLAISYS 在同一输入下生成的 baseline token 序列

最终结果：8 卡生成的 token 序列与单卡严格一致。

这说明以下几个环节整体上是正确的：

• 通信原语调用正确
• 权重切分与加载正确
• Tensor Parallel 前向逻辑正确
• KV Cache 的分片更新与读取正确

4.6 性能观察

分布式部分的性能统计，我将单卡与 8 卡都统一拆成两段：

• load latency：模型初始化、runtime/communicator 初始化、权重加载等阶段
• generate latency：真正进入自回归生成后的阶段

这样做的原因是，如果把 load 和 generate 混在一起统计，就很难判断性能瓶颈究竟来自模型加载还是来自推理阶段的通信/计算。

当前统计结果约为：

• 单卡 load：1.59s
• 单卡 generate：0.51s
• 8 卡 load：6.26s
• 8 卡 generate：4.05s

从这组数据可以看出：

• load latency 变长，主要来自多进程启动、NCCL communicator 初始化和各 rank 加载权重 shard
• generate latency 变长，主要来自 attention 前后的 gather/reduce、额外同步点以及小模型下通信占比过高

结论：

• 当前 8 卡结果以 correctness 验证为主，不以小模型低延迟为目标
• 对 Qwen-1.5B 这类小模型，当前通信开销大于计算收益，因此 8 卡慢于单卡
• 本次工作的主要价值是打通 NCCL 通信、Tensor Parallel 和分布式推理链路，为更大模型或 MoE 扩展提供底座

5. 本次训练营的主要收获

5.1 对 AI 系统分层的理解更具体了

这次我把“构建系统、运行时、算子、模型、Python 包装”真正串了起来，也更清楚每一层在后端接入中的作用。

5.2 对“正确性优先”的工程策略有了更深理解

无论是单卡 CUDA 推理，还是 8 卡 Tensor Parallel，我都优先保证：

• 接口先稳定
• 数学结果先对
• 测试链路先打通
• 再去找真正值得优化的点

5.3 分布式问题的复杂度更多来自系统细节

像 bootstrap 文件、stream 创建顺序、non-contiguous tensor、计时口径这些问题，在真实工程里都非常关键。

6. 结项总结

在 Project #2 中，我完成了 LLAISYS 的 NVIDIA CUDA 后端接入，打通了从构建系统、Runtime、核心算子到 Qwen2 单卡推理的完整流程，并完成了 correctness 和性能验证。

在 Project #5 中，我完成了基于 NCCL 的 8 卡 Tensor Parallel 推理支持，打通了通信层、算子并行层和模型并行层，并验证了 8 卡输出与单卡一致。